package com.ycl.javacore.exam.limit;

import org.junit.Test;

import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * User: 杨成龙
 * Desc: 限流算法
 *
 * 1、令牌桶限流算法
 * 核心概念：
 *  有个任务每秒往桶里放N个token，不能超过最大限制。。。
 *  每次进来执行任务，检查token是否>=1，>=1则token--再执行任务，如果token < 1，则直接拒绝
 *  优点：可以瞬间承载大流量
 *
 * 其余方法：
 *
 * 1、atomicInteger 直接计数，同时有个固定速率的线程每秒把atomicInteger清零
 *
 * 2、SemaphoreOne信号量，获取到信号量则开始执行任务，获取失败，则拒绝。最后释放release。。。同时有个固定速率的线程每秒把信号量全部释放
 */
public class TokenBucket {

    /**
     * 定义了桶
     */
    public class Bucket {
        //容量
        int capacity;
        //速率，每秒放多少
        int rateCount;
        //目前token个数
        AtomicInteger curCount = new AtomicInteger(0);

        public Bucket(int capacity, int rateCount) {
            this.capacity = capacity;
            this.rateCount = rateCount;
        }

        public void put() {
            System.out.println(Thread.currentThread().getName() + "====== put之前的数量：" + curCount.get());
            if (curCount.get() < capacity) {
                System.out.println("目前数量==" + curCount.get() + ", 我还可以继续放");
                curCount.addAndGet(rateCount);
            }
        }

        public boolean get() {
            System.out.println(Thread.currentThread().getName() + "====== get之前的数量：" + curCount.get());
            if (curCount.get() >= 1) {
                curCount.decrementAndGet();
                return true;
            }
            return false;
        }
    }

    @Test
    public void testTokenBucket() throws InterruptedException {

        Bucket bucket = new Bucket(5, 2);

        //固定线程，固定的速率往桶里放数据，比如每秒N个
        ScheduledThreadPoolExecutor scheduledCheck = new ScheduledThreadPoolExecutor(1);
        scheduledCheck.scheduleAtFixedRate(() -> {
            bucket.put();
        }, 0, 1, TimeUnit.SECONDS);

        //先等待一会儿，让桶里放点token
        Thread.sleep(6000);

        //模拟瞬间10个线程进来拿token
        for (int i = 0; i < 10; i++) {
            new Thread(() -> {
                if (bucket.get()) {
                    System.out.println(Thread.currentThread() + "获取到了资源");
                } else {
                    System.out.println(Thread.currentThread() + "被拒绝");
                }
            }).start();
        }

        //等待，往桶里放token
        Thread.sleep(3000);

        //继续瞬间10个线程进来拿token
        for (int i = 0; i < 10; i++) {
            new Thread(() -> {
                if (bucket.get()) {
                    System.out.println(Thread.currentThread() + "获取到了资源");
                } else {
                    System.out.println(Thread.currentThread() + "被拒绝");
                }
            }).start();
        }
    }
}
